﻿// Example 20-02. Using the Mahalanobis distance for classification
#include <algorithm>
#include <ctime>
#include <execution>
#include <iostream>
#include <opencv2/opencv.hpp>

using namespace std;
// using namespace std::execution;  // 简化策略调用

constexpr int    CLUSTER_COUNT = 4;
constexpr int    SAMPLE_COUNT  = 500;
const cv::Scalar colorTab[]    = { cv::Scalar(0, 0, 255),
                                   cv::Scalar(0, 255, 0),
                                   cv::Scalar(255, 0, 0),
                                   cv::Scalar(255, 0, 255),
                                   cv::Scalar(0, 255, 255) };

static void      help(char* argv[])
{
    cout << "\nThis program demonstrates using the Mahalanobis distance for classification.\n"
            " It generates an image with random points, uses kmeans clustering.\n"
            " And then uses the Mahalanobis distance for classification of new points (colors) .\n"
            "Usage:\n"
         << argv[0] << "\n\n"
         << "ESC to quit\n\n"
         << endl;
}

int main(int argc, char** argv)
{
    cv::Mat img(500, 500, CV_8UC3, cv::Scalar::all(0));
    cv::Mat points(SAMPLE_COUNT, 1, CV_32FC2);
    cv::RNG rng(time(nullptr));
    help(argv);
    rng.fill(points, cv::RNG::UNIFORM, cv::Scalar(0, 0), cv::Scalar(img.cols, img.rows));

    cv::Mat labels;
    kmeans(points,
           CLUSTER_COUNT,
           labels,
           cv::TermCriteria(cv::TermCriteria::EPS | cv::TermCriteria::COUNT, 10, 1.0),
           3,
           cv::KMEANS_PP_CENTERS);

    vector<cv::Mat> clusters(CLUSTER_COUNT);

    for (int i = 0; i < SAMPLE_COUNT; i++) {
        int       clusterIdx = labels.at<int>(i);
        cv::Point ipt        = points.at<cv::Point2f>(i);
        cv::Mat   sample(1, 2, CV_32FC1);
        sample.at<float>(0, 0) = ipt.x;
        sample.at<float>(0, 1) = ipt.y;
        clusters[clusterIdx].push_back(sample);
        cv::circle(img, ipt, 2, colorTab[clusterIdx], cv::FILLED, cv::LINE_AA);
    }
    cv::namedWindow("Example 20-02");
    cv::imshow("Example 20-02", img);

    vector<cv::Mat> covarMats(CLUSTER_COUNT);
    vector<cv::Mat> means(CLUSTER_COUNT);
    for (int i = 0; i < CLUSTER_COUNT; i++) {
        cv::calcCovarMatrix(clusters[i], covarMats[i], means[i], cv::COVAR_NORMAL | cv::COVAR_ROWS, 5);
    }

    cout << "Press any button to classify the next point!\n"
         << "Press ESC to exit." << endl;

    for (;;) {
        if (const int key = cv::waitKey(); key == 27) break;

        // 生成随机点
        const cv::Mat newPoint = (cv::Mat_<float>(1, 2) << rng.uniform(0, img.cols), rng.uniform(0, img.rows));

        // 预分配并计算所有马氏距离
        vector<float> mahalanobisDistance(CLUSTER_COUNT);
        std::transform(std::execution::par_unseq,
                       means.begin(),
                       means.end(),
                       covarMats.begin(),
                       mahalanobisDistance.begin(),
                       [&newPoint](const cv::Mat& mean, const cv::Mat& covar) {
                           const auto ret = cv::Mahalanobis(newPoint, mean, covar);
                           return static_cast<float>(ret);
                       });

        // 找到最小距离的聚类
        const auto minIt
            = std::min_element(std::execution::par, mahalanobisDistance.begin(), mahalanobisDistance.end());
        const int clusterIdx = static_cast<int>(distance(mahalanobisDistance.begin(), minIt));

        // 可视化结果
        // const cv::Point2f pt = newPoint.at<cv::Point2f>(0);
        if (auto data = newPoint.ptr<float>(0); data) {
            cout << "New point: (" << data[0] << ", " << data[1] << ")" << endl;
            auto pt = cv::Point2f(data[0], data[1]);
            cv::circle(img, pt, 5, colorTab[clusterIdx], cv::FILLED, cv::LINE_AA);
            cv::imshow("Example 20-02", img);
        }
    }

    cv::destroyAllWindows();
    return 0;
}
